import regex as re
from typing import BinaryIO, Iterable
from tqdm import tqdm
import psutil
from loguru import logger
import os
import time
import json
import copy
import numpy as np


def gpt2_bytes_to_unicode() -> dict[int, str]:
    """
    Returns a mapping between every possible byte (an integer from 0 to 255) to a
    printable unicode string character representation. This function is taken
    from the GPT-2 code.

    For example, `chr(0)` is `\x00`, which is an unprintable character:

    >>> chr(0)
    '\x00'
    >>> print(chr(0))

    As a result, this function returns a dictionary `d` where `d[0]` returns `Ā`.
    The bytes that are visually printable keep their original string representation [1].
    For example, `chr(33)` returns `!`, and so accordingly `d[33]` returns `!`.
    Note in particular that the space character `chr(32)` becomes `d[32]`, which
    returns 'Ġ'.

    For unprintable characters, the function shifts takes the integer representing
    the Unicode code point of that character (returned by the Python `ord`) function
    and shifts it by 256. For example, `ord(" ")` returns `32`, so the the space character
    ' ' is shifted to `256 + 32`. Since `chr(256 + 32)` returns `Ġ`, we use that as the
    string representation of the space.

    This function can simplify the BPE implementation and makes it slightly easier to
    manually inspect the generated merges after they're serialized to a file.
    """
    # These 188 integers can used as-is, since they are not whitespace or control characters.
    # See https://www.ssec.wisc.edu/~tomw/java/unicode.html.
    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
    cs = bs[:]
    # now get the representations of the other 68 integers that do need shifting
    # each will get mapped chr(256 + n), where n will grow from 0...67 in the loop
    # Get printable representations of the remaining integers 68 integers.
    n = 0
    for b in range(2**8):
        if b not in bs:
            # If this integer isn't in our list of visually-representable
            # charcters, then map it to the next nice character (offset by 256)
            bs.append(b)
            cs.append(2**8 + n)
            n += 1
    characters = [chr(n) for n in cs]
    d = dict(zip(bs, characters))
    return d


def find_chunk_boundaries(
    file: BinaryIO, 
    desired_num_chunks: int, 
    split_special_token: bytes
) -> list[int]:
    """
    Chunk the file into parts that can be counted independently.
    May return fewer chunks if the boundaries end up overlapping.
    """
    assert isinstance(split_special_token, bytes), (
        "Must represent special token as a bytestring"
    )

    # Get total file size in bytes
    file.seek(0, os.SEEK_END)
    file_size = file.tell()
    file.seek(0)

    chunk_size = file_size // desired_num_chunks

    # Initial guesses for chunk boundary locations, uniformly spaced
    # Chunks start on previous index, don't include last index
    chunk_boundaries = [i * chunk_size for i in range(desired_num_chunks + 1)]
    chunk_boundaries[-1] = file_size

    mini_chunk_size = 4096  # Read ahead by 4k bytes at a time

    for bi in range(1, len(chunk_boundaries) - 1):
        initial_position = chunk_boundaries[bi]
        file.seek(initial_position)  # Start at boundary guess
        while True:
            mini_chunk = file.read(mini_chunk_size)  # Read a mini chunk

            # If EOF, this boundary should be at the end of the file
            if mini_chunk == b"":
                chunk_boundaries[bi] = file_size
                break

            # Find the special token in the mini chunk
            found_at = mini_chunk.find(split_special_token)
            if found_at != -1:
                chunk_boundaries[bi] = initial_position + found_at
                break
            initial_position += mini_chunk_size

    # Make sure all boundaries are unique, but might be fewer than desired_num_chunks
    return sorted(set(chunk_boundaries))


def load_chunk_data(chunk_bytes: bytes, special_tokens: list[str], _id: int) -> dict[str, int]:
    t0 = time.time()
    logger.info(f"Processor {_id} 开始处理")
    word_freq_dict = {}
    chunk = chunk_bytes.decode('utf-8')
    gpt_pattern = re.compile(r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
    
    if special_tokens:
        special_pattern = f"({'|'.join(re.escape(s) for s in special_tokens)})"
        text_parts = re.split(special_pattern, chunk)
    else:
        text_parts = [chunk]

    for text_part in text_parts:
        if text_part in special_tokens:
            continue

        matches = gpt_pattern.findall(text_part)
        
        for m in matches:
            if len(m.encode('utf-8')) <= 1: continue
            freq = word_freq_dict.get(m, 0)
            word_freq_dict[m] = freq + 1
    t1 = time.time()
    logger.info(f"Processor {_id} 处理结束, 耗时{round(t1 - t0, 2)}s")
    return word_freq_dict


class Tokenizer:
    def __init__(self, vocab: dict[int, bytes], merges: list[tuple[bytes, bytes]], special_tokens: list[str] | None = None):
        
        self.vocab: dict[int, bytes] = vocab
        self.revocab: dict[bytes, int] = {v:k for k,v in vocab.items()}
        self.merges: list[tuple[bytes, bytes]] = merges
        self.special_tokens: list[str] = special_tokens
        if special_tokens:
            self.test_special_tokens()

        self.num_process = 40
        self.gpt_pattern = re.compile(r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
        self.special_pattern = re.compile(f"({'|'.join(re.escape(s) for s in sorted(special_tokens, key=lambda x: len(x), reverse=True))})") if special_tokens else None

        self.merges_dict = {k:i for i, k in enumerate(self.merges)}

    @staticmethod
    def from_files(vocab_filepath: str, merges_filepath: str, special_tokens: list[str] = None) -> "Tokenizer":
        gpt2_byte_decoder = {v: k for k, v in gpt2_bytes_to_unicode().items()}
        with open(vocab_filepath) as vocab_f:
            gpt2_vocab = json.load(vocab_f)
        gpt2_bpe_merges = []
        with open(merges_filepath) as f:
            for line in f:
                cleaned_line = line.rstrip()
                if cleaned_line and len(cleaned_line.split(" ")) == 2:
                    gpt2_bpe_merges.append(tuple(cleaned_line.split(" ")))
        # The GPT-2 tokenizer uses a remapped unicode encoding for bytes. Let's
        # just return the original bytes, so we don't force students to use
        # any particular encoding scheme.
        vocab = {
            int(gpt2_vocab_index): bytes([gpt2_byte_decoder[token] for token in gpt2_vocab_item])
            for gpt2_vocab_index, gpt2_vocab_item in gpt2_vocab.items()
        }
        # If any of the special tokens don't exist in the vocab, append them to the vocab.
        if special_tokens:
            for special_token in special_tokens:
                byte_encoded_special_token = special_token.encode("utf-8")
                if byte_encoded_special_token not in set(vocab.values()):
                    vocab[len(vocab)] = byte_encoded_special_token

        merges = [
            (
                bytes([gpt2_byte_decoder[token] for token in merge_token_1]),
                bytes([gpt2_byte_decoder[token] for token in merge_token_2]),
            )
            for merge_token_1, merge_token_2 in gpt2_bpe_merges
        ]
        return Tokenizer(vocab, merges, special_tokens)

    def pre_tk(self, text: str) -> list[str]:
        res = []
        if self.special_tokens:
            text_parts = self.special_pattern.split(text)
            for text_part in text_parts:
                if text_part in self.special_tokens:
                    res.append(text_part)
                    continue

                matches = self.gpt_pattern.findall(text_part)
                res.extend(matches)
        else:
            matches = self.gpt_pattern.findall(text)
            res.extend(matches)

        return res

    @property
    def vocab_size(self):
        return len(self.revocab)
    
    def encode(self, text: str, return_np: bool = False) -> list[int] | np.ndarray:
        res = []
        pre_texts = self.pre_tk(text)
        for pre_text in pre_texts:
            b_text = pre_text.encode('utf-8')
            token_ids = self.bytes2token(b_text)
            res.extend(token_ids)
        if return_np:
            res = np.array(res, dtype=np.uint16)
        return res
    
    def decode(self, ids: list[int]) -> str:
        b_res = self.token2bytes(ids)
        return b_res.decode('utf-8', errors='ignore')
    
    def encode_iterable(self, iterable: Iterable[str]) -> Iterable[int]:
        for s in iterable:
            res = self.encode(s)
            for r in res:
                yield r

    def token2bytes(self, token_in: list[int] | tuple[int]) -> bytes:
        res = bytearray()
        for t in token_in:
            res.extend(self.vocab[t])
        return bytes(res)
    
    def bytes2token(self, bytes_in: bytes) -> list[int]:
        if bytes_in in self.revocab:
            return [self.revocab[bytes_in]]
        list_bytes = [bytes_in[i:i+1] for i in range(len(bytes_in))]
        
        while True:
            pair_index = []
            for i in range(len(list_bytes) - 1):
                pair = (list_bytes[i], list_bytes[i + 1])
                index = self.merges_dict.get(pair)
                if index:
                    pair_index.append([index, i])

            if not pair_index:
                break

            min_pair_index, list_bytes_index = sorted(pair_index, key=lambda x: x[0])[0]
            pair = self.merges[min_pair_index]
            list_bytes = list_bytes[:list_bytes_index] + [pair[0] + pair[1]] + list_bytes[list_bytes_index + 2:]

        return [self.revocab[byte] for byte in list_bytes]
    

    def add_token(self, token: bytes):
        if token not in self.revocab:
            index = self.vocab_size
            self.vocab[index] = token
            self.revocab[token] = self.vocab_size
        else:
            logger.warning(f'Token ({token}, {token.decode()}) 已经存在了')

    def test_special_tokens(self):
        for token_str in self.special_tokens:
            b_token = token_str.encode('utf-8')
            if b_token not in self.revocab:
                self.revocab[b_token] = self.vocab_size
                self.vocab[self.vocab_size] = b_token


def get_tokenizer_from_vocab_merges_path(
    vocab_path: str | os.PathLike,
    merges_path: str | os.PathLike,
    special_tokens: list[str] | None = None,
):
    gpt2_byte_decoder = {v: k for k, v in gpt2_bytes_to_unicode().items()}
    with open(vocab_path) as vocab_f:
        gpt2_vocab = json.load(vocab_f)
    gpt2_bpe_merges = []
    with open(merges_path) as f:
        for line in f:
            cleaned_line = line.rstrip()
            if cleaned_line and len(cleaned_line.split(" ")) == 2:
                gpt2_bpe_merges.append(tuple(cleaned_line.split(" ")))
    # The GPT-2 tokenizer uses a remapped unicode encoding for bytes. Let's
    # just return the original bytes, so we don't force students to use
    # any particular encoding scheme.
    vocab = {
        gpt2_vocab_index: bytes([gpt2_byte_decoder[token] for token in gpt2_vocab_item])
        for gpt2_vocab_item, gpt2_vocab_index in gpt2_vocab.items()
    }
    # If any of the special tokens don't exist in the vocab, append them to the vocab.
    if special_tokens:
        for special_token in special_tokens:
            byte_encoded_special_token = special_token.encode("utf-8")
            if byte_encoded_special_token not in set(vocab.values()):
                vocab[len(vocab)] = byte_encoded_special_token

    merges = [
        (
            bytes([gpt2_byte_decoder[token] for token in merge_token_1]),
            bytes([gpt2_byte_decoder[token] for token in merge_token_2]),
        )
        for merge_token_1, merge_token_2 in gpt2_bpe_merges
    ]
    return Tokenizer(vocab, merges, special_tokens)

if __name__ == '__main__':
    # import cProfile

    # with cProfile.Profile() as pr:
    tk = get_tokenizer_from_vocab_merges_path(vocab_path='tests/fixtures/gpt2_vocab.json', merges_path='tests/fixtures/gpt2_merges.txt', special_tokens=["<|endoftext|><|endoftext|>", "<|endoftext|>"])

    test_string = 'Hello, how <|endoftext|><|endoftext|> are you?<|endoftext|>'

    encoded_ids = tk.encode(test_string)
    tokenized_string = tk.decode(encoded_ids)
    # Ensure the special <|endoftext|> token is preserved
    assert tokenized_string.count("<|endoftext|>") == 3

    decoded_string = tk.decode(encoded_ids)
    assert test_string == decoded_string